Update app.py

app.py

@@ -13,30 +13,29 @@ from scipy.fftpack import dct
 
 
 
-def crop_image(image_array): #Gradio takes image, and automatically converts into an ndarray.
+def crop_image(image): #Gradio takes image, and automatically converts into an ndarray.
     #This function will return a value which indicates whether the crop succeeded or not
     #and why it failed, if it did. It will also return the cropped image if succeeded, and the original one otherwise.
     #-1 indicates failure because of grayscale or non-RGB.
     #-2 indicates failure because image is too small to crop
     # 1 indicates success.
-    if len(image_array.shape)!=3: #This is a grayscale image.
-        return -1, image_array
-    x_dim, y_dim = image_array.shape[0], image_array.shape[1]
+    if len(np.asarray(image).shape)!=3: #This is a grayscale image.
+        return -1, image
+    x_dim, y_dim = image.size
     if x_dim < 256 or y_dim < 256:
-        return -2, image_array
-    image = im_lib.fromarray(image_array)
+        return -2, image
     left, upper = randint(0, x_dim-256), randint(0, y_dim-256)
     right, lower = 256+left, 256+upper
     image = image.crop((left,upper,right,lower))
-    return 1, np.asarray(image)
+    return 1, image
 
 def HPF_filter(image):
     return im_lib.fromarray(np.asarray(image)-np.asarray(image.filter(ImageFilter.GaussianBlur)))
 
-def final_image_array_single(image_array):
-    cropped_key, image_array = crop_image(image_array)
+def final_image_array_single(image):
+    cropped_key, image = crop_image(image)
     if cropped_key == 1:
-        image_array = np.asarray(HPF_filter(im_lib.fromarray(image_array)))
+        image_array = np.asarray(HPF_filter(image))
     return image_array
 
 def create_model_single(model_weights_file = "single_channel_model_best_val_real_precision_weights"):
@@ -62,29 +61,25 @@ def create_model_single(model_weights_file = "single_channel_model_best_val_real
 
 model_test_single = create_model_single()
 
-def real_or_not_single(image_array, model = model_test_single):
-    cropped_key, cropped_image_array = crop_image(image_array)
+def real_or_not_single(image, model = model_test_single):
+    cropped_key, cropped_image = crop_image(image)
     if cropped_key == -1:
         return "This image cannot be processed as it is not RGB."
     elif cropped_key == -2:
         return "This image is too small to be processed."
     else:
-        image_array = final_image_array_single(cropped_image_array)
+        image_array = final_image_array_single(cropped_image)
         prediction_list = model.predict(image_array.reshape(1,256,256,3))
         if np.argmax(prediction_list) == len(prediction_list[0]) - 1:
             return "This image is probably real."
         else:
-            keywords = ["biggan", "crn", "cyclegan","deepfake","gaugan","imle","progan","san","seeingdark","stargan", "stylegan",
-           "stylegan2","whichfaceisreal"]
-            
-            i = 0
-            while i < len(prediction_list[0]) and prediction_list[0][i] != np.max(prediction_list):
-                i += 1
+            keywords = ["biggan", "crn", "cyclegan","deepfake","gaugan","imle","progan","san","seeingdark","stargan", "stylegan2", "stylegan",
+           "whichfaceisreal"]
+           
 
-            return f"This image is probably fake and generated by {keywords[i]}."
+            return f"This image is probably fake and generated by {keywords[np.argmax(prediction_list)]}."
 
 
-#need to change filters and final image 
 def normalize_image(image,normalizing_factor=255):
     if image.mode == 'RGB':
         return np.asarray(image).reshape(image.size[0],image.size[1],3)/normalizing_factor
@@ -112,10 +107,10 @@ def normalized_gdct(image):
     
 
 def final_image_array_dual(image): 
-    cropped_key, image_array = crop_image(np.array(image))
+    cropped_key, image= crop_image(image)
     if cropped_key == 1:
-        two_images = [normalized_gdct(im_lib.fromarray(image_array)), highpassrgb(im_lib.fromarray(image_array))]
-    return two_images 
+        two_images = [normalized_gdct(image), highpassrgb(image)]
+    return two_images  
 
 def create_model_dual(model_weights_file = "dual_channel_model_best_val_real_precision_weights"):
     try:
@@ -156,32 +151,28 @@ def create_model_dual(model_weights_file = "dual_channel_model_best_val_real_pre
 model_test_dual = create_model_dual()
 
 
-def real_or_not_dual(image_array, model = model_test_dual): #must take an array
-    cropped_key, cropped_image_array = crop_image(image_array)
+def real_or_not_dual(image, model = model_test_dual): #must take an array
+    cropped_key, cropped_image = crop_image(image)
     if cropped_key == -1:
         return "This image cannot be processed as it is not RGB."
     elif cropped_key == -2:
         return "This image is too small to be processed."
     else:
-        image1, image2 = final_image_array_dual(im_lib.fromarray(cropped_image_array))
+        image1, image2 = final_image_array_dual(cropped_image)
         prediction_list = model.predict([image1.reshape(1,256,256,1), image2.reshape(1,256,256,3)])
         if np.argmax(prediction_list) == len(prediction_list[0]) - 1:
             return "This image is probably real."
         else:
-            keywords = ["biggan", "crn", "cyclegan","deepfake","gaugan","imle","progan","san","seeingdark","stargan", "stylegan",
-           "stylegan2","whichfaceisreal"]
+            keywords = ["biggan", "crn", "cyclegan","deepfake","gaugan","imle","progan","san","seeingdark","stargan", "stylegan2", "stylegan",
+           "whichfaceisreal"]
             
-            i = 0
-            while i < len(prediction_list[0]) and prediction_list[0][i] != np.max(prediction_list):
-                i += 1
-
-            return f"This image is probably fake and generated by {keywords[i]}."
 
+            return f"This image is probably fake and generated by {keywords[np.argmax(prediction_list)]}."
 
 
 interface1 = gr.Interface(fn = real_or_not_single,
                          title = "AI or Real?", 
-                         inputs = gr.Image(show_label=False), 
+                         inputs = gr.Image(show_label = False, type = "pil"), 
                          outputs = "text", 
                          description = "<center>Upload your image and we will determine whether it's <strong> real </strong> or <strong> AI-generated </strong> using a Single Channel Neural Network. <br> Please flag any erroneous output.</center>"                
                         )
@@ -189,7 +180,7 @@ interface1 = gr.Interface(fn = real_or_not_single,
 
 interface2 = gr.Interface(fn = real_or_not_dual,
                          title = "AI or Real?", 
-                         inputs = gr.Image(show_label=False), 
+                         inputs = gr.Image(show_label = False, type = "pil"), 
                          outputs = "text", 
                          description = "<center>Upload your image and we will determine whether it's <strong> real </strong> or <strong> AI-generated </strong> using a Dual Channel Neural Network. <br> Please flag any erroneous output.</center>"
                         )