Update app.py

app.py

@@ -20,12 +20,12 @@ with torch.no_grad():
 def guess_drawing(drawing):
     # Access the image data from the 'background' key
     if 'background' in drawing:
-        image_array = np.array(drawing['background'], dtype=np.uint8)  # Extract the background data as an array
+        image_array = np.array(drawing['background'], dtype=np.uint8)
     else:
         return "Invalid drawing format. Unable to process."
 
-    # Convert to PIL image
-    image = Image.fromarray(image_array)
+    # Convert to RGB PIL image to ensure compatibility with CLIP
+    image = Image.fromarray(image_array).convert("RGB")
 
     # Prepare the image for the model
     image_inputs = processor(images=image, return_tensors="pt")
@@ -36,11 +36,17 @@ def guess_drawing(drawing):
 
     # Calculate cosine similarity with each word
     similarity = torch.nn.functional.cosine_similarity(image_features, text_features)
+    
+    # Debug: Print similarity scores for each word
+    for word, score in zip(words, similarity.tolist()):
+        print(f"Similarity score for '{word}': {score}")
+    
     best_match = words[similarity.argmax().item()]
 
     # Return the AI's best guess
     return f"AI's guess: {best_match}"
 
+
 # Set up Gradio interface
 interface = gr.Interface(
     fn=guess_drawing,