Upload 5 files

app.py

@@ -0,0 +1,19 @@
+import gradio as gr
+import numpy as np
+from fish_feeding import FishFeeding
+
+model = FishFeeding()
+model.load_models()
+
+def fish_feeding(images):
+    for i, img in enumerate(images):
+        images[i] = np.array(img, dtype=np.uint8)
+
+    total_feed, times = model.final_fish_feed(images)
+    return {"total_feed": total_feed, "times": times}
+
+inputs = gr.Image(type='numpy', label="Upload fish images")
+outputs = gr.JSON(label="Fish Feeding Results")
+
+app = gr.Interface(fish_feeding, inputs=inputs, outputs=outputs, title="Fish Feeding Predictor")
+app.launch()

counting_model.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c9dc64af07f0ed80fe8f9c9b9d286353136a649accf794c924af6b8832ae07a7
+size 6238297

fish_feeding.py

@@ -0,0 +1,164 @@
+import torch
+import numpy as np
+from PIL import Image
+from transformers import pipeline
+from ultralytics import YOLO
+
+class FishFeeding:
+    
+    def __init__(self, focal_length: float = 27.4) -> None:
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        self.collected_lengths = []
+        self.focal_length = focal_length
+        self.final_weight = None
+        self.length_model_name = "length_model.pt"
+        self.depth_model_name = "vinvino02/glpn-nyu"
+        self.counting_model_name = "counting_model.pt"
+
+    def load_models(self) -> None:
+        self.fish_keypoints_model = YOLO(self.length_model_name)
+        self.depth_model = pipeline(task="depth-estimation", model=self.depth_model_name, device=self.device)
+        self.fish_detection_model = YOLO(self.counting_model_name)
+
+    def predict_fish_length(self, frame):
+        image_obj = Image.fromarray(frame)
+        image_obj = image_obj.resize((640, 640))  # Adjust size as per requirement
+        depth = self.depth_model(image_obj)
+        depth = depth["predicted_depth"]
+        depth = np.array(depth).squeeze()
+
+        results = self.fish_detection_model(frame)[0]
+        if (results.keypoints == None):
+            raise ValueError("No fish detected in the image")
+        keypoints = results.keypoints.xyn[0].detach().cpu().numpy()
+        head = keypoints[0]
+        back = keypoints[1]
+        belly = keypoints[2]
+        tail = keypoints[3]
+
+        depth_w, depth_h = depth.shape[:2]
+
+        head_x = int(head[0] * depth_w)
+        head_y = int(head[1] * depth_h)
+        tail_x = int(tail[0] * depth_w)
+        tail_y = int(tail[1] * depth_h)
+
+        back_x = int(back[0] * depth_w)
+        back_y = int(back[1] * depth_h)
+        belly_x = int(belly[0] * depth_w)
+        belly_y = int(belly[1] * depth_h)
+
+        head_depth = depth[head_y, head_x]
+        tail_depth = depth[tail_y, tail_x]
+
+        fish_length = (
+            np.sqrt(
+                (head_x * head_depth - tail_x * tail_depth) ** 2
+                + (head_y * head_depth - tail_y * tail_depth) ** 2
+            )
+            / self.focal_length
+        )
+        # girth = (
+        #     np.sqrt(
+        #         (back_x * head_depth - belly_x * tail_depth) ** 2
+        #         + (back_y * head_depth - belly_y * tail_depth) ** 2
+        #     )
+        #     / self.focal_length
+        # )
+        return fish_length
+
+    # def videocapture(self):
+    #     cap = cv2.VideoCapture(self.video_path)
+    #     assert cap.isOpened(), "Error reading video file"
+    #     while True:
+    #         ret, frame = cap.read()
+    #         if not ret:
+    #             break
+    #         output = self.predict_fish_length(frame)
+    #         self.collected_lengths.append(output)
+    #     cap.release()
+    #     return self.collected_lengths
+
+    def get_average_weight(self):
+        if not self.collected_lengths:
+            return 0 
+        length_average = sum(self.collected_lengths) / len(self.collected_lengths)
+        final_weight = 0.014 * length_average ** 3.02
+        return final_weight
+
+    def fish_counting(self, images):
+        counting_output = 0
+        for im0 in images:
+            tracks = self.fish_detection_model(im0)
+            counting_output = max(counting_output, len(tracks))
+        
+        return counting_output
+
+    def final_fish_feed(self, images: list):
+        for image in images:
+            try:
+                output = self.predict_fish_length(image)
+            except ValueError:
+                continue
+            self.collected_lengths.append(output)
+
+        average_weight = self.get_average_weight()
+        if 0 <= average_weight <= 50:
+            feed, times = 3.3, 2
+        elif 50 < average_weight <= 100:
+            feed, times = 4.8, 2
+        elif 100 < average_weight <= 250:
+            feed, times = 5.8, 2
+        elif 250 < average_weight <= 500:
+            feed, times = 8.4, 2
+        elif 500 < average_weight <= 750:
+            feed, times = 9.4, 1
+        elif 750 < average_weight <= 1000:
+            feed, times = 10.5, 1
+        elif 1000 < average_weight <= 1500:
+            feed, times = 11.0, 1
+        else: 
+            feed, times = 12.0, 1
+
+        fish_count = self.fish_counting(images)
+        total_feed = feed * fish_count
+        return total_feed, times
+
+
+# if __name__ == "__main__":
+#     to_collect = 6
+#     collected = []
+#     video_path = "object_counting.mp4"
+#     cap = cv2.VideoCapture(video_path)
+
+#     fish_feeding = FishFeeding()
+#     fish_feeding.load_models()
+
+#     d = {"images": []}
+
+#     while True:
+#         ret, frame = cap.read()
+#         if not ret:
+#             break
+
+#         if len(collected) == to_collect:
+#             total_feed, times = fish_feeding.final_fish_feed(collected)
+#             print(f"Total feed: {total_feed}, Feed times: {times}")
+#             collected = []
+
+#             break
+        
+#         collected.append(frame)
+#         d["images"].append(frame.tolist())
+
+#         if cv2.waitKey(1) & 0xFF == ord("q"):
+#             break
+
+#     cap.release()
+#     cv2.destroyAllWindows()
+
+#     # save d to json file
+#     import json
+#     with open("data.json", "w") as f:
+#         json.dump(d, f)
+

length_model.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:208dbde300963cfe83f5f4a9fdfd7a91e640d0410d01ce4e70c96d440cbc03d1
+size 6403287

requirements.txt

@@ -0,0 +1,7 @@
+pillow
+ultralytics
+transformers
+fastapi
+dill==0.3.8
+gradio
+torch