Spaces:
Sleeping
Sleeping
| # Created by yarramsettinaresh GORAKA DIGITAL PRIVATE LIMITED at 01/11/24 | |
| import cv2 | |
| import numpy as np | |
| from ultralytics import YOLO | |
| import torch | |
| import gradio as gr | |
| # Load your models | |
| model = YOLO("chepala_lekka_v5_yolov11n.pt") | |
| # Initialize global video capture variable | |
| cap = None | |
| text_margin = 10 | |
| class_colors = { | |
| "chepa": (255, 0, 0), # Red color for "chepa" | |
| "sanchi": (0, 255, 0), # Green color for "sanchi" | |
| "other_class": (0, 0, 255) # Blue color for other classes | |
| } | |
| class VideoProcessor: | |
| def __init__(self): | |
| self.g_bags = {} | |
| self.g_fishes = {} | |
| self.text_margin = 10 | |
| self.class_colors = { | |
| "chepa": (255, 0, 0), # Red for "chepa" | |
| "sanchi": (0, 255, 0), # Green for "sanchi" | |
| "other_class": (0, 0, 255) # Blue for other classes | |
| } | |
| self.cap = None | |
| def gradio_video_stream(self, video_file): | |
| print(f"gradio_video_stream init : {video_file}") | |
| self.g_bags = {} | |
| self.g_fishes = {} | |
| self.cap = cv2.VideoCapture(video_file) # Open the uploaded video file | |
| while True: | |
| frame = self.process_frame() | |
| if frame is None: | |
| break | |
| yield cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) # Convert BGR to RGB for Gradio | |
| def process_frame(self): | |
| g_fishes = self.g_fishes | |
| g_bags = self.g_bags | |
| cap = self.cap | |
| ret, frame = cap.read() | |
| if not ret: | |
| cap.release() # Release the video capture if no frame is captured | |
| return None # Return None if no frame is captured | |
| frame_height, frame_width = frame.shape[:2] | |
| results = model.track(frame, persist=True) | |
| # person_result = person_model.predict(frame, show=False) | |
| bag_pos = dict() | |
| fishes_pos = dict() | |
| if results[0].boxes.id is not None and results[0].masks is not None: | |
| masks = results[0].masks.data | |
| track_ids = results[0].boxes.id.int().cpu().tolist() | |
| classes = results[0].boxes.cls # Class labels | |
| confidences = results[0].boxes.conf | |
| boxes = results[0].boxes.xyxy | |
| for mask, track_id, cls, conf, box in zip(masks, track_ids, classes, confidences, boxes): | |
| # Convert mask to numpy array if it is a tensor | |
| if isinstance(mask, torch.Tensor): | |
| mask = mask.cpu().numpy() | |
| mask = mask # Convert to numpy array | |
| mask = (mask * 255).astype(np.uint8) # Convert mask to binary format (0 or 255) | |
| # Resize mask to match the original frame dimensions | |
| mask_resized = cv2.resize(mask, (frame_width, frame_height), interpolation=cv2.INTER_NEAREST) | |
| # Get the class name | |
| class_name = model.names[int(cls)] | |
| if class_name == "sanchi": | |
| bag_pos[track_id] = dict(mask=mask) | |
| elif class_name == "chepa": | |
| fishes_pos[track_id] = mask | |
| # Use static color for each class based on the class name | |
| color = class_colors.get(class_name, (255, 255, 255)) # Default to white if class not in color map | |
| color_mask = np.zeros_like(frame, dtype=np.uint8) | |
| color_mask[mask_resized > 128] = color # Apply color where mask is | |
| # Blend original frame with color mask | |
| frame = cv2.addWeighted(frame, 1, color_mask, 0.5, 0) | |
| # Display the label and confidence score on the frame | |
| # Display the label and confidence score on the frame | |
| label = f"{class_name}{track_id}" | |
| position = np.where(mask_resized > 128) | |
| if position[0].size > 0 and position[1].size > 0: | |
| y, x = position[0][0], position[1][0] # Position for label | |
| if not class_name == "sanchi": | |
| cv2.putText(frame, label, (x, y), cv2.FONT_HERSHEY_SIMPLEX, 2, color, 3) | |
| else: | |
| bag_pos[track_id]["xy"] = (x, y) | |
| for fish_id, fish_mask in fishes_pos.items(): | |
| if fish_id not in g_fishes: | |
| g_fishes[fish_id] = dict(in_sanchi=False) | |
| if not g_fishes[fish_id]["in_sanchi"]: | |
| for bag_id, bag_info in bag_pos.items(): | |
| bag_mask = bag_info["mask"] | |
| if np.any(np.logical_and(fish_mask, bag_mask)): | |
| if bag_id not in g_bags: | |
| g_bags[bag_id] = 0 | |
| g_bags[bag_id] += 1 | |
| g_fishes[fish_id]["in_sanchi"] = True | |
| print(g_bags) | |
| for bag_id, v in bag_pos.items(): | |
| color = class_colors.get("sanchi", (255, 255, 255)) | |
| label = f"{g_bags.get(bag_id, 0)}: sanchi{bag_id}" | |
| cv2.putText(frame, label, v["xy"], cv2.FONT_HERSHEY_SIMPLEX, 2, color, 3) | |
| # Loop through each bag position | |
| # Loop through each bag entry | |
| for bag_id, v in g_bags.items(): | |
| if v: | |
| # Set the text color to red | |
| color = (0, 0, 255) # Red color in BGR format | |
| label = f"BAG{bag_id}: {v}" | |
| # Get the size of the text box | |
| (text_width, text_height), baseline = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 2, 5) | |
| # Calculate the position for the rectangle (background) | |
| x1 = frame_width - text_width - text_margin | |
| y1 = text_margin + text_height + baseline | |
| x2 = frame_width - text_margin | |
| y2 = text_margin | |
| # Draw a rectangle for the background | |
| cv2.rectangle(frame, (x1, y2), (x2, y1), (0, 0, 0), thickness=-1) # Black rectangle | |
| # Adjust the transparency (if you still want it) | |
| # Optional: Create an overlay effect | |
| overlay = frame.copy() | |
| cv2.addWeighted(overlay, 0.5, frame, 0.5, 0, frame) # Create transparency effect | |
| # Put the text on top of the rectangle | |
| cv2.putText(frame, label, (x1, y2 + 100), cv2.FONT_HERSHEY_SIMPLEX, 2, color, 3) | |
| return frame | |
| # Gradio interface | |
| iface = gr.Interface(fn=VideoProcessor().gradio_video_stream, | |
| inputs=gr.Video(label="Upload Video"), | |
| outputs=gr.Image(), | |
| ) | |
| iface.launch() | |