Spaces:
Sleeping
Sleeping
File size: 2,984 Bytes
98a3af2 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 |
from abc import ABC, abstractmethod
from pathlib import Path
from PIL import Image
from ultralytics import YOLO
class BaseModel(ABC):
@abstractmethod
def __init__(self):
pass
@abstractmethod
def predict_image(self, image):
pass
@abstractmethod
def predict_video(self, video):
pass
class YOLOModelv1(BaseModel):
"""Model: modelYOLOv8n_datasetDOTAv2_epochs5_batch1.pt"""
def __init__(self):
repo_root = Path(__file__).resolve().parents[1]
weights_path = (
repo_root / "models" / "modelYOLOv8n_datasetDOTAv2_epochs5_batch1.pt"
)
self.model = YOLO(str(weights_path), task="detect")
def predict_image(self, image, min_confidence):
results = self.model.predict(image, save=False, imgsz=640, conf=min_confidence)
annotated_image_filename = "annotated_image.png"
last_im = None
for result in results:
im_array = result.plot()
last_im = Image.fromarray(im_array[..., ::-1]) # RGB PIL image
last_im.save(annotated_image_filename)
# Return PIL Image for robust display in Streamlit
return last_im if last_im is not None else Image.open(annotated_image_filename)
def predict_video(self, video, min_confidence, target_dir_name="annotated_video"):
self.model.predict(
video,
save=True,
project=".",
name=target_dir_name,
exist_ok=True,
imgsz=640,
conf=min_confidence,
)
class YOLOModelv2(BaseModel):
"""Model: modelYOLOv8n_datasetDIOR_epochs50_batch16.pt"""
def __init__(self):
repo_root = Path(__file__).resolve().parents[1]
weights_path = (
repo_root / "models" / "modelYOLOv8n_datasetDIOR_epochs50_batch16.pt"
)
self.model = YOLO(str(weights_path), task="detect")
def predict_image(self, image, min_confidence, classes=None):
results = self.model.predict(
image, save=False, imgsz=800, conf=min_confidence, classes=classes
)
annotated_image_filename = "annotated_image.png"
last_im = None
for result in results:
im_array = result.plot()
last_im = Image.fromarray(im_array[..., ::-1]) # RGB PIL image
last_im.save(annotated_image_filename)
# Return PIL Image for robust display in Streamlit
return last_im if last_im is not None else Image.open(annotated_image_filename)
def predict_video(
self, video, min_confidence, target_dir_name="annotated_video", classes=None
):
self.model.predict(
video,
save=True,
project=".",
name=target_dir_name,
exist_ok=True,
imgsz=800,
conf=min_confidence,
classes=classes,
)
|