takecの気まぐれブログ

プログラミング等の話題を気まぐれに

画像のランキング付け

写真を撮影した際、似たような写真がたくさんあって一番いいのはどれか悩んだことがあった。 ランダムに2枚並べて表示し、どちらの方がより良いか選択していけば良い写真を選択できると思い、 pythonで簡単に作成した。

import sys
from PyQt5.QtCore import Qt
from PyQt5.QtWidgets import (
        QApplication, QWidget, QLabel, QHBoxLayout, QVBoxLayout,
        QGraphicsView, QGraphicsScene, QGraphicsItem
        )
from PyQt5.QtGui import (QPixmap)
import glob
import random
import csv

imgPath = './data/Resize/*.PNG'
dataPath = './result.csv'

class MainWindow(QWidget):
    def __init__(self):
        super().__init__()
        self.title = 'Image rank'
        self.initUI()
        self.left = 50
        self.top = 50
        self.width = 1200
        self.height = 480
        self.setWindowTitle(self.title)
        self.setGeometry(self.left, self.top, self.width, self.height)
        self.getImgList()
        self.setNextPix()
    def initUI(self):
        print("initUI")
        self.view1 = QGraphicsView()
        self.scene1 = QGraphicsScene(self.view1)
        self.view1.setScene(self.scene1)
        self.label1 = QLabel("Label1")

        self.view2 = QGraphicsView()
        self.scene2 = QGraphicsScene(self.view2)
        self.view2.setScene(self.scene2)
        self.label2 = QLabel("Label2")

        self.hlayout = QHBoxLayout()
        self.vlayout1 = QVBoxLayout()
        self.vlayout2 = QVBoxLayout()
        self.vlayout1.addWidget(self.view1)
        self.vlayout1.addWidget(self.label1)
        self.vlayout2.addWidget(self.view2)
        self.vlayout2.addWidget(self.label2)
        self.hlayout.addLayout(self.vlayout1)
        self.hlayout.addLayout(self.vlayout2)
        self.setLayout(self.hlayout)
        self.show()

    def resizeEvent(self, e):
        self.resizeView()

    def resizeView(self):
        self.view1.fitInView(self.scene1.sceneRect(), Qt.KeepAspectRatio)
        self.view2.fitInView(self.scene2.sceneRect(), Qt.KeepAspectRatio)

    def getImgList(self):
        self.imgFiles = glob.glob(imgPath)
        if(len(self.imgFiles)<2):
            print("can't compare")
            exit()
        self.ranking = []
        for i in range(0, len(self.imgFiles)):
            self.ranking.append( [0, self.imgFiles[i]] )
        #print(self.ranking)

    def setNextPix(self):
        while True:
            self.n1 = random.randint(0,len(self.imgFiles)-1)
            self.n2 = random.randint(0,len(self.imgFiles)-1)
            if self.n1 != self.n2:
                break
        #print(self.n1,self.imgFiles[self.n1],self.n2,self.imgFiles[self.n2])
        self.scene1.clear()
        self.pixmap1 = QPixmap(self.imgFiles[self.n1])
        self.scene1.addPixmap(self.pixmap1)
        self.label1.setText(self.imgFiles[self.n1])
        self.scene2.clear()
        self.pixmap2 = QPixmap(self.imgFiles[self.n2])
        self.scene2.addPixmap(self.pixmap2)
        self.label2.setText(self.imgFiles[self.n2])
        self.resizeView()

    def keyPressEvent(self, e):
        if( e.key() == Qt.Key_Q ):
            self.ranking.sort(key=lambda x:x[0], reverse=True)
            with open(dataPath, 'w') as f:
                writer = csv.writer(f, lineterminator='\n')
                writer.writerows(self.ranking)
            self.close()

        if( e.key() == Qt.Key_H or e.key() == Qt.Key_L or e.key() == Qt.Key_M):
            if( e.key() == Qt.Key_H ): # left
                self.ranking[self.n1][0] += 1
            if( e.key() == Qt.Key_L ): # right
                self.ranking[self.n2][0] += 1
            self.setNextPix()


if __name__ == '__main__':
    app = QApplication(sys.argv)
    print("create MainWindow")
    main_window = MainWindow()
    print("exec")
    sys.exit(app.exec_())