Eitan Asher commited on
Commit
4c09ceb
·
1 Parent(s): 05c76ce

Added Gradio app and model files

Browse files
app.py ADDED
@@ -0,0 +1,87 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+ from transformers import ViTForImageClassification, ViTFeatureExtractor, AutoConfig
3
+ import gradio as gr
4
+ from PIL import Image
5
+ import os
6
+ import logging
7
+ from safetensors.torch import load_file # Import safetensors loading function
8
+
9
+ # Set up logging
10
+ logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
11
+
12
+ # Define the directory containing the model files
13
+ model_dir = "." # Use current directory
14
+
15
+ # Define paths to the specific model files
16
+ model_path = os.path.join(model_dir, "model.safetensors")
17
+ config_path = os.path.join(model_dir, "config.json")
18
+ preprocessor_path = os.path.join(model_dir, "preprocessor_config.json")
19
+
20
+ # Check if all required files exist
21
+ for path in [model_path, config_path, preprocessor_path]:
22
+ if not os.path.exists(path):
23
+ logging.error(f"File not found: {path}")
24
+ raise FileNotFoundError(f"Required file not found: {path}")
25
+ else:
26
+ logging.info(f"Found file: {path}")
27
+
28
+ # Load the configuration
29
+ config = AutoConfig.from_pretrained(config_path)
30
+
31
+ # Ensure the labels are consistent with the model's config
32
+ labels = list(config.id2label.values())
33
+ logging.info(f"Labels: {labels}")
34
+
35
+ # Load the feature extractor
36
+ feature_extractor = ViTFeatureExtractor.from_pretrained(preprocessor_path)
37
+
38
+ # Load the model using the safetensors file
39
+ state_dict = load_file(model_path) # Use safetensors to load the model weights
40
+ model = ViTForImageClassification.from_pretrained(
41
+ pretrained_model_name_or_path=None,
42
+ config=config,
43
+ state_dict=state_dict
44
+ )
45
+
46
+ # Ensure the model is in evaluation mode
47
+ model.eval()
48
+ logging.info("Model set to evaluation mode")
49
+
50
+ # Define the prediction function
51
+ def predict(image):
52
+ logging.info("Starting prediction")
53
+ logging.info(f"Input image shape: {image.size}")
54
+
55
+ # Preprocess the image
56
+ logging.info("Preprocessing image")
57
+ inputs = feature_extractor(images=image, return_tensors="pt")
58
+ logging.info(f"Preprocessed input shape: {inputs['pixel_values'].shape}")
59
+
60
+ logging.info("Running inference")
61
+ with torch.no_grad():
62
+ outputs = model(**inputs)
63
+ logits = outputs.logits
64
+ probabilities = torch.nn.functional.softmax(logits[0], dim=0)
65
+
66
+ logging.info(f"Raw logits: {logits}")
67
+ logging.info(f"Probabilities: {probabilities}")
68
+
69
+ # Prepare the output dictionary
70
+ result = {labels[i]: float(probabilities[i]) for i in range(len(labels))}
71
+ logging.info(f"Prediction result: {result}")
72
+
73
+ return result
74
+
75
+ # Set up the Gradio Interface
76
+ logging.info("Setting up Gradio interface")
77
+ gradio_app = gr.Interface(
78
+ fn=predict,
79
+ inputs=gr.Image(type="pil"),
80
+ outputs=gr.Label(num_top_classes=6),
81
+ title="Dress Length Classifier"
82
+ )
83
+
84
+ # Launch the app
85
+ if __name__ == "__main__":
86
+ logging.info("Launching the app")
87
+ gradio_app.launch()
models/.DS_Store ADDED
Binary file (6.15 kB). View file
 
models/best_model_/config.json ADDED
@@ -0,0 +1,42 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "_name_or_path": "google/vit-base-patch16-224-in21k",
3
+ "architectures": [
4
+ "ViTForImageClassification"
5
+ ],
6
+ "attention_probs_dropout_prob": 0.0,
7
+ "encoder_stride": 16,
8
+ "hidden_act": "gelu",
9
+ "hidden_dropout_prob": 0.0,
10
+ "hidden_size": 768,
11
+ "id2label": {
12
+ "0": "Mini",
13
+ "1": "Midi",
14
+ "2": "Maxi",
15
+ "3": "Knee",
16
+ "4": "Floor",
17
+ "5": "Ankle",
18
+ "6": "Asymmetric"
19
+ },
20
+ "image_size": 224,
21
+ "initializer_range": 0.02,
22
+ "intermediate_size": 3072,
23
+ "label2id": {
24
+ "Ankle": 5,
25
+ "Asymmetric": 6,
26
+ "Floor": 4,
27
+ "Knee": 3,
28
+ "Maxi": 2,
29
+ "Midi": 1,
30
+ "Mini": 0
31
+ },
32
+ "layer_norm_eps": 1e-12,
33
+ "model_type": "vit",
34
+ "num_attention_heads": 12,
35
+ "num_channels": 3,
36
+ "num_hidden_layers": 12,
37
+ "patch_size": 16,
38
+ "problem_type": "single_label_classification",
39
+ "qkv_bias": true,
40
+ "torch_dtype": "float32",
41
+ "transformers_version": "4.47.1"
42
+ }
models/best_model_/model.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:b0fb6b8053ee72176edc1fbe5a3b74bf2d8df14ac3829506366c302b6fefdc21
3
+ size 343239356
models/best_model_/preprocessor_config.json ADDED
@@ -0,0 +1,23 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "do_convert_rgb": null,
3
+ "do_normalize": true,
4
+ "do_rescale": true,
5
+ "do_resize": true,
6
+ "image_mean": [
7
+ 0.5,
8
+ 0.5,
9
+ 0.5
10
+ ],
11
+ "image_processor_type": "ViTFeatureExtractor",
12
+ "image_std": [
13
+ 0.5,
14
+ 0.5,
15
+ 0.5
16
+ ],
17
+ "resample": 2,
18
+ "rescale_factor": 0.00392156862745098,
19
+ "size": {
20
+ "height": 224,
21
+ "width": 224
22
+ }
23
+ }
models/best_model_/requirements.txt ADDED
@@ -0,0 +1,4 @@
 
 
 
 
 
1
+ torch
2
+ transformers
3
+ gradio
4
+ pillow
models/best_model_/training_args.bin ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:785f162c2592baf2f7144442f212ba5cc0aadf93584f59eb314c2c61da6b2732
3
+ size 5496