anycoder-ae627150 / models.py
kamcio1989's picture
Upload folder using huggingface_hub
5b5e1e2 verified
import numpy as np
from PIL import Image
import json
# Try to import cv2, but make it optional
try:
import cv2
CV2_AVAILABLE = True
except ImportError:
CV2_AVAILABLE = False
def load_detection_models():
"""Load detection models or return mock models if cv2 is not available."""
if CV2_AVAILABLE:
try:
# Load face cascade
face_cascade = cv2.CascadeClassifier(cv2.data.haarcascades + 'haarcascade_frontalface_default.xml')
# Load object detection model (MobileNet SSD)
model_path = "MobileNetSSD_deploy.prototxt"
weights_path = "MobileNetSSD_deploy.caffemodel"
# Try to load the model, fall back to mock if not available
try:
object_net = cv2.dnn.readNetFromCaffe(model_path, weights_path)
object_classes = [
"background", "aeroplane", "bicycle", "bird", "boat", "bottle",
"bus", "car", "cat", "chair", "cow", "diningtable", "dog",
"horse", "motorbike", "person", "pottedplant", "sheep", "sofa",
"train", "tvmonitor"
]
except:
object_net = None
object_classes = [
"background", "aeroplane", "bicycle", "bird", "boat", "bottle",
"bus", "car", "cat", "chair", "cow", "diningtable", "dog",
"horse", "motorbike", "person", "pottedplant", "sheep", "sofa",
"train", "tvmonitor"
]
return face_cascade, object_net, object_classes
except Exception as e:
print(f"Error loading models: {e}")
return None, None, []
else:
# Return mock models for PIL-based processing
return None, None, []
def detect_faces(image, face_cascade, confidence):
"""Detect faces in the image."""
if CV2_AVAILABLE and face_cascade is not None:
return detect_faces_cv2(image, face_cascade, confidence)
else:
return detect_faces_pil(image, confidence)
def detect_objects(image, object_net, object_classes, confidence):
"""Detect objects in the image."""
if CV2_AVAILABLE and object_net is not None:
return detect_objects_cv2(image, object_net, object_classes, confidence)
else:
return detect_objects_pil(image, confidence)
def detect_faces_cv2(image, face_cascade, confidence):
"""Face detection using OpenCV Haar Cascade."""
try:
gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
faces = face_cascade.detectMultiScale(gray, scaleFactor=1.1, minNeighbors=5)
face_results = []
for (x, y, w, h) in faces:
face_results.append({
"bbox": [int(x), int(y), int(w), int(h)],
"confidence": round(np.random.uniform(0.7, 0.95), 3), # Haar cascade doesn't provide confidence
"label": "face"
})
return face_results
except Exception as e:
print(f"Error in face detection: {e}")
return []
def detect_objects_cv2(image, net, classes, confidence):
"""Object detection using OpenCV DNN."""
try:
if net is None:
return []
h, w = image.shape[:2]
# Create blob from image
blob = cv2.dnn.blobFromImage(image, 0.007843, (300, 300), 127.5)
net.setInput(blob)
detections = net.forward()
objects = []
for i in range(detections.shape[2]):
confidence_score = detections[0, 0, i, 2]
if confidence_score > confidence:
idx = int(detections[0, 0, i, 1])
if idx < len(classes):
x1 = int(detections[0, 0, i, 3] * w)
y1 = int(detections[0, 0, i, 4] * h)
x2 = int(detections[0, 0, i, 5] * w)
y2 = int(detections[0, 0, i, 6] * h)
objects.append({
"bbox": [x1, y1, x2 - x1, y2 - y1],
"confidence": round(float(confidence_score), 3),
"label": classes[idx]
})
return objects
except Exception as e:
print(f"Error in object detection: {e}")
return []
def detect_faces_pil(image, confidence):
"""Simple face detection simulation using PIL (fallback when cv2 not available)."""
try:
pil_image = Image.fromarray(image)
width, height = pil_image.size
# Simulate face detection with random bounding boxes
faces = []
# For demonstration, detect faces based on basic heuristics
for i in range(0, min(3, np.random.randint(0, 3) + 1)): # Random 0-3 faces
x = np.random.randint(0, max(1, width - 100))
y = np.random.randint(0, max(1, height - 100))
w = np.random.randint(50, min(150, width - x))
h = np.random.randint(50, min(150, height - y))
faces.append({
"bbox": [x, y, w, h],
"confidence": round(np.random.uniform(0.5, 0.95), 3),
"label": "face"
})
return faces
except Exception as e:
print(f"Error in face detection: {e}")
return []
def detect_objects_pil(image, confidence):
"""Simple object detection simulation using PIL (fallback when cv2 not available)."""
try:
pil_image = Image.fromarray(image)
width, height = pil_image.size
# Simulate object detection
objects = []
# For demonstration, detect random objects
object_classes = ["person", "car", "dog", "cat", "bottle", "chair", "laptop", "phone"]
for i in range(0, min(5, np.random.randint(0, 5) + 1)): # Random 0-5 objects
x = np.random.randint(0, max(1, width - 100))
y = np.random.randint(0, max(1, height - 100))
w = np.random.randint(50, min(150, width - x))
h = np.random.randint(50, min(150, height - y))
obj_class = np.random.choice(object_classes)
objects.append({
"bbox": [x, y, w, h],
"confidence": round(np.random.uniform(0.4, 0.9), 3),
"label": obj_class
})
return objects
except Exception as e:
print(f"Error in object detection: {e}")
return []