Update app.py

app.py

@@ -9,48 +9,47 @@ import firebase_admin
 from firebase_admin import credentials, firestore, storage
 from pydantic import BaseModel
 
-# Load the pre-trained model
-learn = load_learner('model.pkl')
-
-# Define categories and map them to indices
-searches = ['formal', 'casual', 'athletic']
-searches = sorted(searches)  # Ensure the categories are in sorted order
-values = [i for i in range(0, len(searches))]
-class_dict = dict(zip(searches, values))
-
-# Set up logging
-logging.basicConfig(level=logging.DEBUG,
-                    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
-logger = logging.getLogger(__name__)
-
-# Initialize Firebase
-try:
-    cred = credentials.Certificate("serviceAccountKey.json")
-    firebase_app = firebase_admin.initialize_app(cred, {
-        'storageBucket': 'future-forge-60d3f.appspot.com'
-    })
-    db = firestore.client()
-    bucket = storage.bucket(app=firebase_app)
-    logger.info("Firebase initialized successfully")
-except Exception as e:
-    logger.error(f"Failed to initialize Firebase: {str(e)}")
-
-app = FastAPI()
-
-# Add CORS middleware
-app.add_middleware(
-    CORSMiddleware,
-    allow_origins=["*"],
-    allow_credentials=True,
-    allow_methods=["*"],
-    allow_headers=["*"],
-)
-
-
-# Define the input model
-class FileProcess(BaseModel):
-    file_path: str
+import torch
+from transformers import AutoImageProcessor, AutoModelForObjectDetection
+from PIL import Image, ImageDraw, ImageFont
+import cv2
+import random
 
+# Load model and processor
+processor = AutoImageProcessor.from_pretrained("valentinafeve/yolos-fashionpedia")
+model = AutoModelForObjectDetection.from_pretrained("valentinafeve/yolos-fashionpedia")
+
+# Fashionpedia categories
+FASHION_CATEGORIES = [
+    'shirt, blouse', 'top, t-shirt, sweatshirt', 'sweater', 'cardigan', 'jacket', 'vest', 'pants', 'shorts', 'skirt',
+    'coat', 'dress', 'jumpsuit', 'cape', 'glasses', 'hat', 'headband, head covering, hair accessory', 'tie', 'glove',
+    'watch', 'belt', 'leg warmer', 'tights, stockings', 'sock', 'shoe', 'bag, wallet', 'scarf', 'umbrella', 'hood',
+    'collar', 'lapel', 'epaulette', 'sleeve', 'pocket', 'neckline', 'buckle', 'zipper', 'applique', 'bead', 'bow',
+    'flower', 'fringe', 'ribbon', 'rivet', 'ruffle', 'sequin', 'tassel'
+]
+
+
+def detect_fashion(image):
+    inputs = processor(images=image, return_tensors="pt")
+    outputs = model(**inputs)
+
+    # Convert outputs (bounding boxes and class logits) to COCO API
+    target_sizes = torch.tensor([image.size[::-1]])
+    results = processor.post_process_object_detection(outputs, threshold=0.1, target_sizes=target_sizes)[0]
+
+    detected_items = []
+    for score, label, box in zip(results["scores"], results["labels"], results["boxes"]):
+        if score > 0.5:  # Adjust this threshold as needed
+            detected_items.append((FASHION_CATEGORIES[label], score.item(), box.tolist()))
+
+    return detected_items
+
+
+def check_dress_code(detected_items):
+    formal_workplace_attire = {
+        "shirt, blouse", "jacket", "tie", "coat", "sweater", "cardigan", "coat"
+    }
+    return any(item[0] in formal_workplace_attire for item in detected_items)
 
 @app.post("/process")
 async def process_file(file_data: FileProcess):
@@ -70,9 +69,9 @@ async def process_file(file_data: FileProcess):
         file_type = file_data.file_path.split('.')[-1].lower()
 
         try:
-            if file_type in ['jpg', 'jpeg', 'png', 'bmp']:
-                output = process_video(str(tmp_file_path))
-                result = {"type": "image", "data": {"result": output}}
+            if file_type in ['mp4', 'avi', 'mov', 'wmv']:
+                output,testing = process_video(str(tmp_file_path))
+                result = {"type": "video", "data": {"result": output}}
             else:
                 raise HTTPException(status_code=400, detail="Unsupported file type")
 
@@ -96,22 +95,27 @@ async def process_file(file_data: FileProcess):
         raise HTTPException(status_code=500, detail=f"Error processing file: {str(e)}")
 
 
-def process_video(video_path):
-    # Load the image from the provided path
-    img = PILImage.create(video_path)
-    
-    # Make the prediction
-    classification, _, probs = learn.predict(img)
-    
-    # Convert the prediction to a confidence dictionary
-    confidences = {label: float(probs[i]) for i, label in enumerate(class_dict)}
-    
-    # If classification is not formal, return 'informal'
-    if classification != 'formal':
-        informal_confidence = sum(confidences[label] for label in class_dict if label != 'formal')
-        return {'informal': informal_confidence}
-    else:
-        return {'formal': confidences['formal']}
+def process_video(video_path,num_frames=10):
+    cap = cv2.VideoCapture(video_path)
+    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
+
+    frame_indices = sorted(random.sample(range(total_frames), min(num_frames, total_frames)))
+
+    compliance_results = []
+
+    for frame_index in frame_indices:
+        cap.set(cv2.CAP_PROP_POS_FRAMES, frame_index)
+        ret, frame = cap.read()
+        if ret:
+            image = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
+            detected_items = detect_fashion(image)
+            is_compliant = check_dress_code(detected_items)
+            compliance_results.append(is_compliant)
+
+    cap.release()
+
+    average_compliance = sum(compliance_results) / len(compliance_results)
+    return average_compliance, compliance_results
 
 if __name__ == "__main__":
     logger.info("Starting the Face Emotion Recognition API")