adi-123 commited on
Commit
44f0770
·
verified ·
1 Parent(s): 3289c11

Upload 3 files

Browse files
Files changed (3) hide show
  1. app.py +57 -0
  2. best.pt +3 -0
  3. requirements.txt +7 -0
app.py ADDED
@@ -0,0 +1,57 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+ import streamlit as st
3
+ import cv2
4
+ import numpy as np
5
+ from ultralytics import YOLO
6
+ from PIL import Image
7
+ import tempfile
8
+
9
+ # Directly set the path for the model
10
+ MODEL_PATH = 'best.pt'
11
+
12
+ # Initialize YOLO model with custom trained weights
13
+ model = YOLO(MODEL_PATH)
14
+
15
+ def detect_rhino_image(image):
16
+ image = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
17
+ results = model(image)[0]
18
+ for box in results.boxes.data.tolist():
19
+ x1, y1, x2, y2, score, class_id = box
20
+ if score > 0.5:
21
+ cv2.rectangle(image, (int(x1), int(y1)), (int(x2), int(y2)), (0, 255, 0), 4)
22
+ cv2.putText(image, results.names[int(class_id)].upper(), (int(x1), int(y1 - 10)), cv2.FONT_HERSHEY_SIMPLEX, 1.3, (0, 255, 0), 3, cv2.LINE_AA)
23
+ return image
24
+
25
+ def detect_rhino_video(video_file):
26
+ cap = cv2.VideoCapture(video_file.name)
27
+ ret, frame = cap.read()
28
+ H, W, _ = frame.shape
29
+ out = cv2.VideoWriter(video_file.name + '_output.mp4', cv2.VideoWriter_fourcc(*'MP4V'), int(cap.get(cv2.CAP_PROP_FPS)), (W, H))
30
+ while ret:
31
+ results = model(frame)[0]
32
+ for box in results.boxes.data.tolist():
33
+ x1, y1, x2, y2, score, class_id = box
34
+ if score > 0.5:
35
+ cv2.rectangle(frame, (int(x1), int(y1)), (int(x2), int(y2)), (0, 255, 0), 4)
36
+ cv2.putText(frame, results.names[int(class_id)].upper(), (int(x1), int(y1 - 10)), cv2.FONT_HERSHEY_SIMPLEX, 1.3, (0, 255, 0), 3, cv2.LINE_AA)
37
+ out.write(frame)
38
+ ret, frame = cap.read()
39
+ cap.release()
40
+ out.release()
41
+ cv2.destroyAllWindows()
42
+ return video_file.name + '_output.mp4'
43
+
44
+ st.title('Rhinoceros Detection App')
45
+
46
+ st.write("Upload an image or video of rhinoceroses for detection.")
47
+
48
+ file = st.file_uploader("Choose a file...", type=["jpg", "jpeg", "png", "mp4"])
49
+ if file is not None:
50
+ if file.type.split('/')[0] == 'image':
51
+ image = Image.open(file)
52
+ st.image(detect_rhino_image(image), caption='Processed Image', use_column_width=True)
53
+ elif file.type.split('/')[0] == 'video':
54
+ tfile = tempfile.NamedTemporaryFile(delete=False)
55
+ tfile.write(file.read())
56
+ processed_video = detect_rhino_video(tfile)
57
+ st.video(processed_video)
best.pt ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:6c4b614dd7c3d31f93433a971217e7f9cbff05fbd2b1db3cf12f530ce529689a
3
+ size 22539865
requirements.txt ADDED
@@ -0,0 +1,7 @@
 
 
 
 
 
 
 
 
1
+ streamlit
2
+ opencv-python-headless
3
+ numpy
4
+ Pillow
5
+ torch
6
+ torchvision
7
+ ultralytics