Spaces:
Sleeping
Sleeping
YarramsettiNaresh
commited on
Commit
•
1e2d295
1
Parent(s):
93f3bed
init
Browse files- app.py +154 -0
- chepala_lekka_v5_yolov11n.pt +3 -0
- 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
|