Moditha24 commited on
Commit
4a80f4f
Β·
verified Β·
1 Parent(s): e98a614

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +95 -0
app.py ADDED
@@ -0,0 +1,95 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ import torch
3
+ import torch.nn as nn
4
+ import torch.nn.functional as F
5
+ import torchvision.transforms as transforms
6
+ from PIL import Image
7
+ from ResNet_for_CC import CC_model # Import the model
8
+
9
+ # Set device (CPU/GPU)
10
+ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
11
+
12
+ # Load the trained CC_model
13
+ model_path = "CC_net.pt"
14
+ model = CC_model(num_classes1=14)
15
+
16
+ # Load model weights
17
+ state_dict = torch.load(model_path, map_location=device)
18
+ model.load_state_dict(state_dict, strict=False)
19
+ model.to(device)
20
+ model.eval()
21
+
22
+ # Clothing1M Class Labels
23
+ class_labels = [
24
+ "T-Shirt", "Shirt", "Knitwear", "Chiffon", "Sweater", "Hoodie",
25
+ "Windbreaker", "Jacket", "Downcoat", "Suit", "Shawl", "Dress",
26
+ "Vest", "Underwear"
27
+ ]
28
+
29
+ # βœ… **Updated Image Preprocessing Function**
30
+ def preprocess_image(image):
31
+ """Applies necessary transformations to the input image."""
32
+ transform = transforms.Compose([
33
+ transforms.Resize(256),
34
+ transforms.CenterCrop(224),
35
+ transforms.ToTensor(),
36
+ transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
37
+ ])
38
+ return transform(image).unsqueeze(0).to(device)
39
+
40
+ # βœ… **Classification Function**
41
+ def classify_image(image):
42
+ """Processes the input image and returns the predicted clothing category."""
43
+ print("\n[INFO] Received image for classification.")
44
+
45
+ try:
46
+ image = Image.fromarray(image) # Ensure conversion to PIL format
47
+ image = preprocess_image(image) # Apply transformations
48
+ print("[INFO] Image transformed and moved to device.")
49
+
50
+ with torch.no_grad():
51
+ output = model(image)
52
+
53
+ # βœ… Ensure output is a tensor (handle tuple case)
54
+ if isinstance(output, tuple):
55
+ output = output[1] # Extract the actual output tensor
56
+
57
+ print(f"[DEBUG] Model output shape: {output.shape}")
58
+ print(f"[DEBUG] Model output values: {output}")
59
+
60
+ if output.shape[1] != 14:
61
+ return f"[ERROR] Model output mismatch! Expected 14 but got {output.shape[1]}."
62
+
63
+ # Convert logits to probabilities
64
+ probabilities = F.softmax(output, dim=1)
65
+ print(f"[DEBUG] Softmax probabilities: {probabilities}")
66
+
67
+ # Get predicted class index
68
+ predicted_class = torch.argmax(probabilities, dim=1).item()
69
+ print(f"[INFO] Predicted class index: {predicted_class} (Class: {class_labels[predicted_class]})")
70
+
71
+ # Validate and return the prediction
72
+ if 0 <= predicted_class < len(class_labels):
73
+ predicted_label = class_labels[predicted_class]
74
+ confidence = probabilities[0][predicted_class].item() * 100
75
+ return f"Predicted Class: {predicted_label} (Confidence: {confidence:.2f}%)"
76
+ else:
77
+ return "[ERROR] Model returned an invalid class index."
78
+
79
+ except Exception as e:
80
+ print(f"[ERROR] Exception during classification: {e}")
81
+ return "Error in classification. Check console for details."
82
+
83
+ # βœ… **Gradio Interface**
84
+ interface = gr.Interface(
85
+ fn=classify_image,
86
+ inputs=gr.Image(type="numpy"),
87
+ outputs="text",
88
+ title="Clothing1M Image Classifier",
89
+ description="Upload a clothing image, and the model will classify it into one of the 14 categories."
90
+ )
91
+
92
+ # βœ… **Run the Interface**
93
+ if __name__ == "__main__":
94
+ print("[INFO] Launching Gradio interface...")
95
+ interface.launch()