download and clean

dataCollectionFilter.py

import numpy as np
import tensorflow as tf
import cv2 as cv
from pathlib import Path
from send2trash import send2trash
from google_images_download import google_images_download



def downloadImages(keyword):
    response = google_images_download.googleimagesdownload()
    arguments = { "keywords":keyword, "limit":50, "format":'jpg', "size":"medium", "print_urls": False,
                 "image_directory":'./'+keyword.lower()+'/'}
    response.download(arguments)
    #also download png format images
    arguments['format'] = "png"
    response.download(arguments)
    return


ssdLiteMobileNetClasses = """
unclassified
person
bicycle
car
motorcycle
airplane
bus
train
truck
boat
traffic
fire
street
stop
parking
bench
bird
cat
dog
horse
sheep
cow
elephant
bear
zebra
giraffe
hat
backpack
umbrella
shoe
eye
handbag
tie
suitcase
frisbee
skis
snowboard
sports
kite
baseball
baseball
skateboard
surfboard
tennis
bottle
plate
wine
cup
fork
knife
spoon
bowl
banana
apple
sandwich
orange
broccoli
carrot
hot
pizza
donut
cake
chair
couch
potted
bed
mirror
dining
window
desk
toilet
door
tv
laptop
mouse
remote
keyboard
cell
microwave
oven
toaster
sink
refrigerator
blender
book
clock
vase
scissors
teddy
hair
toothbrush
hair
banner
blanket
branch
bridge
building-other
bush
cabinet
cage
cardboard
carpet
ceiling-other
ceiling-tile
cloth
clothes
clouds
counter
cupboard
curtain
desk-stuff
dirt
door-stuff
fence
floor-marble
floor-other
floor-stone
floor-tile
floor-wood
flower
fog
food-other
fruit
furniture-other
grass
gravel
ground-other
hill
house
leaves
light
mat
metal
mirror-stuff
moss
mountain
mud
napkin
net
paper
pavement
pillow
plant-other
plastic
platform
playingfield
railing
railroad
river
road
rock
roof
rug
salad
sand
sea
shelf
sky-other
skyscraper
snow
solid-other
stairs
stone
straw
structural-other
table
tent
textile-other
towel
tree
vegetable
wall-brick
wall-concrete
wall-other
wall-panel
wall-stone
wall-tile
wall-wood
water-other
waterdrops
window-blind
window-other
wood
"""
ssdLiteMobileNetClassesList = ssdLiteMobileNetClasses.split()

keywordList = [
    {
        "name": "chair",
        "ssdLiteMobileNetClassId" : 62
    },
    {
        "name": "bed",
        "ssdLiteMobileNetClassId": 65
    },
    {
        "name": "table",
        "ssdLiteMobileNetClassId": 165
    },
    {
        "name": "cabinet",
        "ssdLiteMobileNetClassId": 98
    },
    {
        "name": "bench",
        "ssdLiteMobileNetClassId" : 15
    },
    {
        "name": "couch",
        "ssdLiteMobileNetClassId": 63
    }
]

def cropAndFilterImages(itemName, graph_def, folder):
    with tf.compat.v1.Session() as sess:
        # Restore session
        sess.graph.as_default()
        tf.import_graph_def(graph_def, name='')

        files = list(Path(folder).glob('*.jpg'))
        files.extend(list(Path(folder).glob('*.png')))
        #ssd mobilenet classes
        """
            0 => unclassified 1 =>person 2 => bicycle 3 => car 4 => motorcycle 5 => airplane
            6 => bus 7 => train 8 => truck 9 => boat 10 => traffic 11 => fire
            12 => street 13 => stop 14 => parking 15 => bench . . . . 
            62 => chair 63 => couch 64 => potted 65 => bed 98 => cabinet 165 => table
        """
                
        #furnitures = [15, 62, 63, 65, 98, 165]
        ssdCurrentClassId = ssdLiteMobileNetClassesList.index(itemName)
        print("ssdClassId: "+str(ssdCurrentClassId) + " for item name: "+itemName)

        for fid,fname in enumerate(files): 
            # Read and preprocess an image.
            img = cv.imread(str(fname))
            if(img is not None):
                rows = img.shape[0]
                cols = img.shape[1]
                inp = cv.resize(img, (300, 300))
                inp = inp[:, :, [2, 1, 0]]  # BGR2RGB

                # Run the model
                out = sess.run([sess.graph.get_tensor_by_name('num_detections:0'),
                                sess.graph.get_tensor_by_name('detection_scores:0'),
                                sess.graph.get_tensor_by_name('detection_boxes:0'),
                                sess.graph.get_tensor_by_name('detection_classes:0')],
                               feed_dict={'image_tensor:0': inp.reshape(1, inp.shape[0], inp.shape[1], 3)})

                # Visualize detected bounding boxes.
                num_detections = int(out[0][0])
                #print(fname, num_detections)
                for i in range(num_detections):
                    classId = int(out[3][0][i])
                    
                    #if(classId in furnitures):
                    if classId == ssdCurrentClassId:
                        score = float(out[1][0][i])
                        bbox = [float(v) for v in out[2][0][i]]
                        if score > 0.3:
                            x = int(bbox[1] * cols)
                            y = int(bbox[0] * rows)
                            right = int(bbox[3] * cols)
                            bottom = int(bbox[2] * rows)
                            cv.imwrite(folder+str(fid)+'-'+str(i)+'.png', img[y:bottom, x:right])

            try:
                printf("trashing: "+str(fname))
                send2trash(str(fname))
            except: pass    
        print('cleaning '+itemName+' done.')
    return    



# Read the graph.
with tf.io.gfile.GFile('./ssdlite_mobilenet_v2_coco_2018_05_09/frozen_inference_graph.pb', 'rb') as f:
    graph_def = tf.compat.v1.GraphDef()
    graph_def.ParseFromString(f.read())
    
for item in keywordList:
    keyword = item["name"]
    downloadImages(keyword)
    cropAndFilterImages(keyword, graph_def, 'downloads/'+keyword.lower()+'/')