Spaces:
Sleeping
Sleeping
| import torch | |
| import numpy as np | |
| import gradio as gr | |
| from PIL import Image | |
| import multiprocessing | |
| import tensorflow as tf | |
| from RealESRGAN import RealESRGAN | |
| from tensorflow.keras import layers | |
| from tensorflow.keras.models import Model | |
| from tensorflow.keras.optimizers import Adam | |
| from tensorflow.keras.applications import InceptionV3 | |
| # function to return model | |
| def create_model(): | |
| inp_shape = (200,200,3) | |
| base_model = InceptionV3(input_shape=inp_shape, include_top=False, weights='imagenet') | |
| x = layers.Flatten()(base_model.output) | |
| x = layers.Dense(256, activation='relu')(x) | |
| x = layers.Dropout(0.5)(x) | |
| output = layers.Dense(8, activation='softmax')(x) | |
| clf_model = Model(inputs=base_model.input, outputs=output) | |
| clf_model.compile(optimizer=Adam(learning_rate=0.0001), | |
| loss='categorical_crossentropy', | |
| metrics=['accuracy']) | |
| return clf_model | |
| clf_model = create_model() | |
| clf_model.load_weights('modelac90.weights.h5') | |
| # input image preprocessing and prediction | |
| device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') | |
| esrgan_model = RealESRGAN(device, scale=4) | |
| esrgan_model.load_weights('weights/RealESRGAN_x4.pth') | |
| enhancement_model = esrgan_model | |
| def classify_logo(inp_image): | |
| pil_img = Image.fromarray(inp_image).resize((200,200), resample=0) | |
| image = np.array(pil_img).astype(np.float16)/255.0 | |
| new_img = np.expand_dims(image, axis=0) | |
| predictions = clf_model.predict(new_img) | |
| labels = ['Adidas Fake', 'Adidas Real', 'Allen Solly Fake', 'Allen Solly Real', | |
| 'Puma Fake', 'Puma Real', 'Us Polo Fake', 'Us Polo Real'] | |
| pred_dict = {} | |
| for i in range(len(labels)): | |
| pred_dict[labels[i]] = predictions[0][i] | |
| return pred_dict | |
| enhaced_images_pipe = [] | |
| def enhance_image(input_image): | |
| try: | |
| enhanced_image = enhancement_model.predict(input_image).convert('RGB') | |
| global enhanced_images_pipe | |
| enhaced_images_pipe.append(np.array(enhanced_image)) | |
| return | |
| except Exception as e: | |
| return None | |
| def fake_logo_detection(input_image, flag): | |
| print("Input image shape => ", input_image.shape) | |
| print("flag => ", flag) | |
| if flag == "Enhance Image before prediction (Uses Real ESRGAN 4X)": | |
| process = multiprocessing.Process(target=enhance_image, args=(input_image,)) | |
| process.start() | |
| process.join(500) | |
| if process.is_alive(): | |
| process.terminate() | |
| process.join() | |
| return "Enhancement model prediction taking too long time⌛. Please! try without enhancement." | |
| else: | |
| global enhanced_images_pipe | |
| if len(enhaced_images_pipe) != 0: | |
| input_image = enhaced_images_pipe[-1] | |
| arr = enhaced_images_pipe.pop() | |
| else: | |
| return "An error occurred while enhancing the image 😞." | |
| return classify_logo(input_image) | |
| # --------------gradio interface------------------------- | |
| input_choices = gr.Dropdown(["No enhancement (prefer for Medium-High resolution images)", "Enhance Image before prediction (Uses Real ESRGAN 4X)"], | |
| value="No enhancement (prefer for Medium-High resolution images)", label="Choose Option") | |
| input_image = gr.Image(sources=['upload'], label="Input Image") | |
| iface = gr.Interface( | |
| fn = fake_logo_detection, | |
| inputs = [input_image, input_choices], | |
| outputs=gr.Label(num_top_classes=2, label="Confidence percentages"), | |
| title="Fake Logo Detection", | |
| description="Upload an image to check if it contains a fake or original logo." | |
| ) | |
| if __name__ == '__main__': | |
| iface.launch(debug=False) |