YarramsettiNaresh commited on
Commit
1e2d295
1 Parent(s): 93f3bed
Files changed (3) hide show
  1. app.py +154 -0
  2. chepala_lekka_v5_yolov11n.pt +3 -0
  3. requirements.txt +3 -0
app.py ADDED
@@ -0,0 +1,154 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Created by yarramsettinaresh GORAKA DIGITAL PRIVATE LIMITED at 01/11/24
2
+ import cv2
3
+ import numpy as np
4
+ from ultralytics import YOLO
5
+ import torch
6
+ import gradio as gr
7
+
8
+ # Load your models
9
+ model = YOLO("chepala_lekka_v5_yolov11n.pt")
10
+
11
+ # Initialize global video capture variable
12
+ cap = None
13
+ text_margin = 10
14
+ class_colors = {
15
+ "chepa": (255, 0, 0), # Red color for "chepa"
16
+ "sanchi": (0, 255, 0), # Green color for "sanchi"
17
+ "other_class": (0, 0, 255) # Blue color for other classes
18
+ }
19
+
20
+ class VideoProcessor:
21
+ def __init__(self):
22
+ self.g_bags = {}
23
+ self.g_fishes = {}
24
+ self.text_margin = 10
25
+ self.class_colors = {
26
+ "chepa": (255, 0, 0), # Red for "chepa"
27
+ "sanchi": (0, 255, 0), # Green for "sanchi"
28
+ "other_class": (0, 0, 255) # Blue for other classes
29
+ }
30
+ self.cap = None
31
+
32
+ def gradio_video_stream(self, video_file):
33
+ print(f"gradio_video_stream init : {video_file}")
34
+ self.g_bags = {}
35
+ self.g_fishes = {}
36
+ self.cap = cv2.VideoCapture(video_file.name) # Open the uploaded video file
37
+
38
+ while True:
39
+ frame = self.process_frame()
40
+ if frame is None:
41
+ break
42
+ yield cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) # Convert BGR to RGB for Gradio
43
+
44
+ def process_frame(self):
45
+ g_fishes = self.g_fishes
46
+ g_bags = self.g_bags
47
+ cap = self.cap
48
+ ret, frame = cap.read()
49
+ if not ret:
50
+ cap.release() # Release the video capture if no frame is captured
51
+ return None # Return None if no frame is captured
52
+
53
+ frame_height, frame_width = frame.shape[:2]
54
+
55
+ results = model.track(frame, persist=True)
56
+ # person_result = person_model.predict(frame, show=False)
57
+ bag_pos = dict()
58
+ fishes_pos = dict()
59
+ if results[0].boxes.id is not None and results[0].masks is not None:
60
+ masks = results[0].masks.data
61
+ track_ids = results[0].boxes.id.int().cpu().tolist()
62
+ classes = results[0].boxes.cls # Class labels
63
+ confidences = results[0].boxes.conf
64
+ boxes = results[0].boxes.xyxy
65
+
66
+ for mask, track_id, cls, conf, box in zip(masks, track_ids, classes, confidences, boxes):
67
+ # Convert mask to numpy array if it is a tensor
68
+ if isinstance(mask, torch.Tensor):
69
+ mask = mask.cpu().numpy()
70
+ mask = mask # Convert to numpy array
71
+ mask = (mask * 255).astype(np.uint8) # Convert mask to binary format (0 or 255)
72
+
73
+ # Resize mask to match the original frame dimensions
74
+ mask_resized = cv2.resize(mask, (frame_width, frame_height), interpolation=cv2.INTER_NEAREST)
75
+
76
+ # Get the class name
77
+ class_name = model.names[int(cls)]
78
+ if class_name == "sanchi":
79
+ bag_pos[track_id] = dict(mask=mask)
80
+ elif class_name == "chepa":
81
+ fishes_pos[track_id] = mask
82
+ # Use static color for each class based on the class name
83
+ color = class_colors.get(class_name, (255, 255, 255)) # Default to white if class not in color map
84
+ color_mask = np.zeros_like(frame, dtype=np.uint8)
85
+ color_mask[mask_resized > 128] = color # Apply color where mask is
86
+
87
+ # Blend original frame with color mask
88
+ frame = cv2.addWeighted(frame, 1, color_mask, 0.5, 0)
89
+
90
+ # Display the label and confidence score on the frame
91
+ # Display the label and confidence score on the frame
92
+ label = f"{class_name}{track_id}"
93
+ position = np.where(mask_resized > 128)
94
+ if position[0].size > 0 and position[1].size > 0:
95
+
96
+ y, x = position[0][0], position[1][0] # Position for label
97
+
98
+ if not class_name == "sanchi":
99
+ cv2.putText(frame, label, (x, y), cv2.FONT_HERSHEY_SIMPLEX, 2, color, 3)
100
+ else:
101
+ bag_pos[track_id]["xy"] = (x, y)
102
+ for fish_id, fish_mask in fishes_pos.items():
103
+ if fish_id not in g_fishes:
104
+ g_fishes[fish_id] = dict(in_sanchi=False)
105
+ if not g_fishes[fish_id]["in_sanchi"]:
106
+ for bag_id, bag_info in bag_pos.items():
107
+ bag_mask = bag_info["mask"]
108
+ if np.any(np.logical_and(fish_mask, bag_mask)):
109
+ if bag_id not in g_bags:
110
+ g_bags[bag_id] = 0
111
+ g_bags[bag_id] += 1
112
+ g_fishes[fish_id]["in_sanchi"] = True
113
+ print(g_bags)
114
+ for bag_id, v in bag_pos.items():
115
+ color = class_colors.get("sanchi", (255, 255, 255))
116
+ label = f"{g_bags.get(bag_id, 0)}: sanchi{bag_id}"
117
+ cv2.putText(frame, label, v["xy"], cv2.FONT_HERSHEY_SIMPLEX, 2, color, 3)
118
+ # Loop through each bag position
119
+ # Loop through each bag entry
120
+ for bag_id, v in g_bags.items():
121
+ if v:
122
+ # Set the text color to red
123
+ color = (0, 0, 255) # Red color in BGR format
124
+ label = f"BAG{bag_id}: {v}"
125
+
126
+ # Get the size of the text box
127
+ (text_width, text_height), baseline = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 2, 5)
128
+
129
+ # Calculate the position for the rectangle (background)
130
+ x1 = frame_width - text_width - text_margin
131
+ y1 = text_margin + text_height + baseline
132
+ x2 = frame_width - text_margin
133
+ y2 = text_margin
134
+
135
+ # Draw a rectangle for the background
136
+ cv2.rectangle(frame, (x1, y2), (x2, y1), (0, 0, 0), thickness=-1) # Black rectangle
137
+
138
+ # Adjust the transparency (if you still want it)
139
+ # Optional: Create an overlay effect
140
+ overlay = frame.copy()
141
+ cv2.addWeighted(overlay, 0.5, frame, 0.5, 0, frame) # Create transparency effect
142
+
143
+ # Put the text on top of the rectangle
144
+ cv2.putText(frame, label, (x1, y2 + 100), cv2.FONT_HERSHEY_SIMPLEX, 2, color, 3)
145
+
146
+ return frame
147
+
148
+ # Gradio interface
149
+ iface = gr.Interface(fn=VideoProcessor().gradio_video_stream,
150
+ inputs=gr.File(label="Upload Video"),
151
+ outputs=gr.Image(),
152
+ )
153
+
154
+ iface.launch()
chepala_lekka_v5_yolov11n.pt ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:ab54db0a8b591c51067aed2a724b133ccb119a715279b09b6eb36c3c61bc1980
3
+ size 6012957
requirements.txt ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ gradio
2
+ opencv-python
3
+ ultralytics