Spaces:
Sleeping
Sleeping
| import numpy as np | |
| import torch | |
| import torch.nn.functional as F | |
| from torchvision.transforms.functional import normalize | |
| import gradio as gr | |
| from gradio_imageslider import ImageSlider | |
| from briarmbg import BriaRMBG | |
| import PIL | |
| from PIL import Image, ImageDraw, ImageFont | |
| from typing import Tuple | |
| net = BriaRMBG.from_pretrained("briaai/RMBG-1.4") | |
| device = torch.device("cuda" if torch.cuda.is_available() else "cpu") | |
| net.to(device) | |
| def draw_ruler(image, actual_height_in_inches = 13, dpi=300, unit="in", color="black"): | |
| # Load the plant image in its original size | |
| canvas = image | |
| plant_width, plant_height = canvas.size | |
| #actual_height_in_inches = 13 #only do the changes in this and generate the images ,,,,,, | |
| #think about approximate height of plant suppose it is 8.3 inch so mention here 9 .and run the code | |
| # Calculate pixels per inch based on the plant image's pixel height | |
| pixels_per_inch = plant_height / actual_height_in_inches | |
| # Create a blank canvas larger than the plant image to accommodate rulers | |
| canvas_width = plant_width #+ 150 # 50 pixels extra for the left ruler | |
| canvas_height = plant_height #+ 150 # 50 pixels extra for the bottom ruler | |
| #canvas = Image.new('RGB', (canvas_width, canvas_height), 'white') | |
| draw = ImageDraw.Draw(canvas) | |
| font = ImageFont.load_default(size=60) | |
| # for drawing scale value on vertical side ................ 0, 0.5, 1, 1.5, 2, 2.5, 3, 3.5...................so on | |
| for i in range(0, int(plant_height / (pixels_per_inch * 0.5)) + 1): | |
| # Calculate position and inch value | |
| #y_position = canvas_height - 150 - int(i * pixels_per_inch * 0.5) | |
| y_position = canvas_height - 0 - int(i * pixels_per_inch * 0.5) | |
| inch_value = i * 0.5 | |
| # Draw tick mark and add text label every 1 inch for better readability | |
| draw.line((0, y_position, 10, y_position), fill='black') | |
| if inch_value.is_integer() and inch_value % 5 == 0: # Display labels only at whole numbers | |
| draw.text((15, y_position - 5), f"{int(inch_value)} in", fill='black',font=font) | |
| #else: | |
| # draw.text((15, y_position - 5), f"{inch_value:.1f} in", fill='black',font=font) | |
| # for drawing scale value on horizontal side ................ 0, 0.5, 1, 1.5, 2, 2.5, 3, 3.5...................so on | |
| for i in range(0, int(plant_width / (pixels_per_inch * 0.5)) + 1): | |
| # Calculate position and inch value | |
| #x_position = 150 + int(i * pixels_per_inch * 0.5) | |
| x_position = 0 + int(i * pixels_per_inch * 0.5) | |
| inch_value = i * 0.5 | |
| # Draw tick mark and add text label every 1 inch for better readability | |
| draw.line((x_position, canvas_height - 40, x_position, canvas_height), fill='black') | |
| if inch_value.is_integer() and inch_value % 5 == 0: # Display labels only at whole numbers | |
| draw.text((x_position, canvas_height - 65), f"{int(inch_value)} in", fill='black',font=font) | |
| #else: | |
| #draw.text((x_position, canvas_height - 25), f"{inch_value:.1f} in", fill='black',font=font) | |
| # Position to paste the plant image on the canvas (leaving space for rulers) | |
| #plant_position = (150, canvas_height - plant_height - 150) | |
| # Paste the original plant image onto the canvas without resizing | |
| #canvas.paste(plant_image, plant_position) | |
| # Save the final image with the plant and rulers | |
| #canvas.save('image_with_plant_and_rulers_high_quality_ONE.png') | |
| #print("The image with the plant and rulers has been saved at original quality.") | |
| return canvas; | |
| def draw_ruler0(image, dpi=300, unit="in", color="black"): | |
| # Define the pot dimensions and distance from camera | |
| pot_diameter_inches = 4 | |
| pot_height_inches = 6 | |
| distance_from_camera_feet = 0.5 | |
| zoom_level = 1 # Zoom level | |
| # Convert distance from feet to inches | |
| distance_from_camera_inches = distance_from_camera_feet * 12 | |
| # Calculate the scale factor (assuming a simple pinhole camera model) | |
| # Scale factor = (actual height of the object) / (height of the object in the image) | |
| # Adjust for zoom level | |
| scale_factor = (pot_height_inches / distance_from_camera_inches) * zoom_level | |
| print ("scale factor:",scale_factor) | |
| # Draw the scale on the image | |
| draw = ImageDraw.Draw(image) | |
| font = ImageFont.load_default() | |
| # Define the position and size of the scale | |
| scale_length_pixels = int(pot_height_inches / scale_factor) | |
| scale_position = (50, 50) # Adjust as needed | |
| # Draw the scale line | |
| draw.line((scale_position[0], scale_position[1], scale_position[0], scale_position[1] + scale_length_pixels), fill="red", width=5) | |
| print("line:",(scale_position[0], scale_position[1], scale_position[0], scale_position[1] + scale_length_pixels)) | |
| # Add text label for the scale | |
| draw.text((scale_position[0] + 10, scale_position[1] + scale_length_pixels // 2), f"{pot_height_inches} inches", fill="red", font=font) | |
| return image | |
| def draw_ruler1(image, dpi=300, unit="in", color="black"): | |
| """ | |
| Draws a ruler on the top and left edges of an image. | |
| Args: | |
| image: The PIL Image object to draw on. | |
| dpi: The DPI of the image (default 300). | |
| unit: The unit of measurement (default "in"). | |
| color: The color of the ruler (default "black"). | |
| """ | |
| #dpi = image.info.get("dpi", (72, 72)) | |
| #dpi = dpi[0] if isinstance(dpi, tuple) else dpi | |
| # Define the pot dimensions and distance from camera | |
| pot_diameter_inches = 4 | |
| pot_height_inches = 6 | |
| distance_from_camera_feet = 3 | |
| zoom_level = 0.5 # Zoom level | |
| # Convert distance from feet to inches | |
| distance_from_camera_inches = distance_from_camera_feet * 12 | |
| # Calculate the scale factor (assuming a simple pinhole camera model) | |
| # Scale factor = (actual height of the object) / (height of the object in the image) | |
| # Adjust for zoom level | |
| scale_factor = (pot_height_inches / distance_from_camera_inches) * zoom_level | |
| scale_length_pixels = int(pot_height_inches / scale_factor) | |
| print ("scale factor:",scale_factor,"scale_length_pixels:",scale_length_pixels) | |
| draw = ImageDraw.Draw(image) | |
| width, height = image.size | |
| font = ImageFont.load_default(size=60) | |
| # Draw top ruler | |
| #for i in range(0, width, int(dpi)): | |
| for i in range(0, width, scale_length_pixels): | |
| x = i | |
| draw.line((x, 0, x, 10), fill=color) | |
| if i % int(scale_length_pixels) == 0: | |
| if scale_length_pixels< 70 : | |
| draw.text((x, 12), str(i // scale_length_pixels), fill=color, font=font) | |
| if scale_length_pixels > 70 and i % 5 == 0: | |
| draw.text((x, 12), str(i // scale_length_pixels), fill=color, font=font) | |
| # Draw left ruler | |
| for i in range(0, height, int(dpi)): | |
| y = i | |
| draw.line((0, y, 10, y), fill=color) | |
| if i % int(dpi) == 0: | |
| if i != 0: | |
| draw.text((12, y), str(i // dpi), fill=color, font=font) | |
| draw.text((100, y-60), "inch", fill=color, font=font) | |
| return image | |
| def mm_to_pixels(mm, dpi): | |
| if dpi is None: | |
| dpi = 72 # Default DPI value | |
| if mm is None: | |
| mm = 0 # Default mm value | |
| return int(mm * dpi / 25.4) | |
| def crop_image(image, left_mm, top_mm, right_mm, bottom_mm): | |
| # Open the image file | |
| #img = Image.open(image) | |
| img = image | |
| print("img.width:",img.width,"img.height:",img.height) | |
| # Get the image DPI (dots per inch) | |
| #dpi = img.info.get('dpi', (72, 72))[0] # Default to 72 DPI if not available | |
| # Get the DPI value as a single number | |
| dpi = img.info.get("dpi", (72, 72)) | |
| dpi = dpi[0] if isinstance(dpi, tuple) else dpi | |
| print ("dpi:",dpi) | |
| # Convert mm to pixels | |
| left = mm_to_pixels(left_mm, dpi) | |
| top = mm_to_pixels(top_mm, dpi) | |
| right = mm_to_pixels(right_mm, dpi) | |
| bottom = mm_to_pixels(bottom_mm, dpi) | |
| print("left:",left,"top:",top,"right:",right,"bottom:",bottom) | |
| if img.width - right <= left : | |
| right=img.width-left-1 | |
| # Crop the image | |
| cropped_img = img.crop((left, top, img.width - right, img.height - bottom)) | |
| return cropped_img | |
| def resize_image(image): | |
| image = image.convert('RGB') | |
| model_input_size = (1024, 1024) | |
| image = image.resize(model_input_size, Image.BILINEAR) | |
| return image | |
| def process(image,left_mm, top_mm, right_mm, bottom_mm, ruler, resize,actual_height_in_inches,background): | |
| # prepare input | |
| orig_image = Image.fromarray(image) | |
| print("orig size:",orig_image.size) | |
| w,h = orig_im_size = orig_image.size | |
| image = resize_image(orig_image) | |
| new_im=image; | |
| if background: | |
| # remove background | |
| im_np = np.array(image) | |
| im_tensor = torch.tensor(im_np, dtype=torch.float32).permute(2,0,1) | |
| im_tensor = torch.unsqueeze(im_tensor,0) | |
| im_tensor = torch.divide(im_tensor,255.0) | |
| im_tensor = normalize(im_tensor,[0.5,0.5,0.5],[1.0,1.0,1.0]) | |
| if torch.cuda.is_available(): | |
| im_tensor=im_tensor.cuda() | |
| #inference | |
| result=net(im_tensor) | |
| # post process | |
| result = torch.squeeze(F.interpolate(result[0][0], size=(h,w), mode='bilinear') ,0) | |
| ma = torch.max(result) | |
| mi = torch.min(result) | |
| result = (result-mi)/(ma-mi) | |
| # image to pil | |
| im_array = (result*255).cpu().data.numpy().astype(np.uint8) | |
| pil_im = Image.fromarray(np.squeeze(im_array)) | |
| # paste the mask on the original image | |
| new_im = Image.new("RGBA", pil_im.size, (0,0,0,0)) | |
| new_im.paste(orig_image, mask=pil_im) | |
| # new_orig_image = orig_image.convert('RGBA') | |
| else: | |
| new_im=image; | |
| new_im= crop_image(new_im, left_mm, top_mm, right_mm, bottom_mm) | |
| if ruler: | |
| new_im = draw_ruler(new_im,actual_height_in_inches) | |
| if resize is not None and resize > 0: | |
| new_im = new_im.resize( [int(resize * s) for s in new_im.size],Image.Resampling.LANCZOS ) | |
| print("resize:",new_im.size) | |
| return new_im | |
| # return [new_orig_image, new_im] | |
| # block = gr.Blocks().queue() | |
| # with block: | |
| # gr.Markdown("## BRIA RMBG 1.4") | |
| # gr.HTML(''' | |
| # <p style="margin-bottom: 10px; font-size: 94%"> | |
| # This is a demo for BRIA RMBG 1.4 that using | |
| # <a href="https://huggingface.co/briaai/RMBG-1.4" target="_blank">BRIA RMBG-1.4 image matting model</a> as backbone. | |
| # </p> | |
| # ''') | |
| # with gr.Row(): | |
| # with gr.Column(): | |
| # input_image = gr.Image(sources=None, type="pil") # None for upload, ctrl+v and webcam | |
| # # input_image = gr.Image(sources=None, type="numpy") # None for upload, ctrl+v and webcam | |
| # run_button = gr.Button(value="Run") | |
| # with gr.Column(): | |
| # result_gallery = gr.Gallery(label='Output', show_label=False, elem_id="gallery", columns=[1], height='auto') | |
| # ips = [input_image] | |
| # run_button.click(fn=process, inputs=ips, outputs=[result_gallery]) | |
| # block.launch(debug = True) | |
| # block = gr.Blocks().queue() | |
| gr.Markdown("## BRIA RMBG 1.4") | |
| gr.HTML(''' | |
| <p style="margin-bottom: 10px; font-size: 94%"> | |
| This app is leveraging the RMBG v1.4 model which is developed on the IS-Net enhanced with unique training scheme and proprietary dataset. | |
| These modifications significantly improve the model’s accuracy and effectiveness in diverse image-processing scenarios. | |
| This is an internal tool being tested for rootsraja.in product images at | |
| <a href="https://huggingface.co/spaces/laitkor/product-image" target="_blank">Rootsraja internal product tool using an image matting model</a> as backbone. | |
| </p> | |
| ''') | |
| title = "Background Removal, Cropping, Resizing, Ruler and scaling" | |
| description = r"""<a href="https://huggingface.co/spaces/laitkor/product-image" target="_blank">Rootsraja internal product tool<br> | |
| To use upload your image and wait. Read more at model card <a href='https://huggingface.co/briaai/RMBG-1.4' target='_blank'><b>briaai/RMBG-1.4</b></a>.<br> | |
| """ | |
| examples = [['./input.jpg'],] | |
| # output = ImageSlider(position=0.5,label='Image without background', type="pil", show_download_button=True) | |
| # demo = gr.Interface(fn=process,inputs="image", outputs=output, examples=examples, title=title, description=description) | |
| demo = gr.Interface(fn=process, | |
| inputs=[ | |
| gr.Image(type="numpy", label="Upload Image"), | |
| gr.Number(label="Left Crop (mm)",value=0), | |
| gr.Number(label="Top Crop (mm)",value=0), | |
| gr.Number(label="Right Crop (mm)",value=0), | |
| gr.Number(label="Bottom Crop (mm)",value=0), | |
| gr.Checkbox(label="Ruler!"), | |
| gr.Number(label="Resize (between 0.1 and 0.9)",value=0), | |
| gr.Number(label="Plant height in inches",value=0), | |
| gr.Checkbox(label="Remove Background?") | |
| ], | |
| outputs="image", examples=examples, title=title, description=description) | |
| if __name__ == "__main__": | |
| demo.launch(share=False) |